from typing import Dict, Iterator
from app.jobs.pdfparser.internal.logic.splitter.text_splitter import TextSegment, TextSplitter
from langchain_text_splitters import RecursiveCharacterTextSplitter


class TextParagraphSplitter(TextSplitter):
    """按段落分割文本"""
    def __init__(self, paragraph_separator: str = "\n\n", segment_size: int = 800) -> None:
        super().__init__()
        self._paragraph_separator = paragraph_separator
        self._segment_size = segment_size
        self._recursive_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self._segment_size,
            chunk_overlap=80,
            length_function=len,
            separators=[
                        ".", "。", "\u3002",
                        ",", "，", "\uff0c",
                        "\n",
                        " ", "\u3000", "\u200b",    # 空格/全角空格/零宽空格
                        ''],
            is_separator_regex=False,
        )
        
    def split(self, text: str, extra_kwargs: Dict = None) -> Iterator[TextSegment]:
        if not (isinstance(text, str) and text):
            raise ValueError("Input text must be a non-empty string")
        
        # 第一阶段：按段落分割
        paragraphs = (p.strip() for p in text.split(self._paragraph_separator))
        filtered_paragraphs = (p for p in paragraphs if p)
        
        # 第二阶段：对超长段落进一步进行分割
        for index, paragraph in enumerate(filtered_paragraphs, start=1):
            if len(paragraph) <= self._segment_size:
                yield TextSegment(
                    text=paragraph,
                    meta={
                        "paragraph_index": index,
                        "segment_index": 1,
                        "is_recursive": False
                    },
                    splitter_id=self.splitter_id
                )
            else:
                for segment_index, segment in enumerate(
                    self._recursive_splitter.split_text(paragraph), start=1
                ):
                    yield TextSegment(
                        text=segment,
                        meta={
                            "paragraph_index": index,
                            "segment_index": segment_index,
                            "is_recursive": True,
                            "original_length": len(paragraph)
                        },
                        splitter_id=self.splitter_id
                    )
